import numpy as np
import cv2
import traceback
import os

import apriltag
from camera_utils.camera import camera_thread, read_yaml
from camera_utils.config import CALIB_DIR

# apriltag参数
with open('april_tag_size', 'r')as f:
    apriltag_size = {}
    l_name = ["w", "h", "left", "right", "up", "down", "z"]
    ss = f.read().split('\n')
    for s in ss[1:-1]:
        res = s.split(' ')
        order = res[0]
        res = res[1:]
        apriltag_size[order] = {}
        for l, item in zip(l_name, res):
            apriltag_size[order][l] = float(item)


# 双相机采集图像
def two_camera_grub():
    # two camera
    caps = [[camera_thread(i), read_yaml(i)[1:3]] for i in range(2)]

    try:
        # apriltag info
        order = "08"

        w = apriltag_size[order]["w"]
        h = apriltag_size[order]["h"]
        left = apriltag_size[order]["left"]
        right = apriltag_size[order]["right"]
        down = apriltag_size[order]["down"]
        up = apriltag_size[order]["up"]

        frames = []  # 图像
        calis = []  # 标定信息
        # two frame
        for cap, cali in caps:
            r, frame = cap.read()
            calis.append(cali)
            frames.append(frame)

        # two camera's window
        cv2.namedWindow("AprilTag0", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("AprilTag0", 500, 300)
        cv2.moveWindow("AprilTag0", 100, 100)
        cv2.namedWindow("AprilTag1", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("AprilTag1", 500, 300)
        cv2.moveWindow("AprilTag1", 600, 300)

        # 创建一个apriltag检测器
        at_detector = apriltag.Apriltag()
        at_detector.create_detector()

        # 预显示画面，根据apriltag显示调整顺序
        # 绘制顺序颜色，红蓝绿浅蓝
        colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (255, 255, 0)]
        while cv2.waitKey(1) != ord('q'):
            # 进行apriltag检测
            for im in frames:
                tag = at_detector.detect(im)
                if tag.shape[0] == 4:
                    for i in range(4):
                        cv2.circle(im, tuple(tag[i]), 2, colors[i], -1)

            cv2.imshow("AprilTag0", frames[0])
            cv2.imshow("AprilTag1", frames[1])

            k = 0
            for cap, _ in caps:
                r, frame = cap.read()
                frames[k] = frame.copy()
                k += 1

        point_order = input(
            "Please input order:")  # should be like 1234 2341 (1 means the right_down,2 means the left_down)
        # 内部点
        object_points_sample = np.array([
            [left + w, down, 0],  # right_down
            [left, down, 0],  # left_down
            [left, down + h, 0],  # left_top
            [left + w, down + h, 0]  # right_top
        ],

            dtype=np.float64)
        # 外部点
        object_points_2_sample = np.array([[left + w + right, 0, 0],  # right_down
                                           [0, 0, 0],  # left_down
                                           [0, down + h + up, 0],  # left_top
                                           [left + w + right, down + h + up, 0]],  # right_top
                                          dtype=np.float64)

        object_points = object_points_sample.copy()
        object_points_2 = object_points_2_sample.copy()
        if len(point_order):
            assert len(point_order) == 4
            for i in range(4):
                object_points[point_order.find(str(i))] = object_points_sample[i]
                object_points_2[point_order.find(str(i))] = object_points_2_sample[i]

        stop_criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
                         30, 0.001)

        locate_flag = True

        j = 0

        if os.path.exists(CALIB_DIR):
            ts = os.listdir(CALIB_DIR)
            if len(ts):
                j = np.argmax(np.array([int(t.split('_')[1]) for t in ts])) + 1
        else:
            os.mkdir(CALIB_DIR)

        key = cv2.waitKey(1)
        while key != ord('q') & 0xFF:
            k = 0
            cam_r = True
            for cap, _ in caps:
                r, frame = cap.read()
                cam_r = cam_r & r
                frames[k] = frame.copy()
                k += 1
            if not cam_r:
                break
            detect_flag = True
            if locate_flag:
                for i, data in enumerate(zip(frames, calis)):
                    frame, cali = data
                    tag = at_detector.detect(frame)

                    if tag.shape[0] != 4:
                        # detect failure, to the next frame
                        detect_flag = False
                        break

                    _, R, tvec = cv2.solvePnP(object_points, np.float64(tag).reshape(-1, 2), cali[0], cali[1])

                    p = cv2.projectPoints(object_points_2, cv2.Rodrigues(R)[0], tvec, cali[0], cali[1])[0].astype(
                        np.float32).reshape(-1, 1, 2)

                    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                    p = cv2.cornerSubPix(gray, p, (5, 5), (-1, -1), stop_criteria).reshape(-1, 2)

                    im = frame.copy()

                    for i in range(4):
                        cv2.circle(frame, tuple(p[i].astype(int)), 4, colors[i], 2)

                    cv2.imshow("AprilTag{0}".format(i), frame)
                    key = cv2.waitKey(0)
                    # TODO: 键位设定，c保存，其他键skip
                    if key == ord('c') & 0xFF:
                        # 存储标定图片，前一张为没有标记标定点的，后一张有以便查看
                        cv2.imwrite(f"/{i}_c{j}.jpg", im)

                        cv2.imwrite(f"{CALIB_DIR}/{i}_c{j}_p.jpg", frame)

                        # f.write("%d\n" % j)
                        with open(f"{CALIB_DIR}/{i}_{j}_corner.txt", "w") as f:
                            for ip in p:
                                f.write("%0.09f %0.09f \n" % (ip[0], ip[1]))

                        j += 1
                    else:
                        print("throw")

                locate_flag = False
            else:
                for i, frame in enumerate(frames):
                    cv2.imshow("AprilTag{0}".format(i), frame)

            if not detect_flag:
                # 如果检测失败，就直接显示图片
                for i, frame in enumerate(frames):
                    cv2.putText(frame, "detect_failure", (100, 100), cv2.FONT_HERSHEY_COMPLEX, 3, (0, 0, 255), 1)
                    cv2.imshow("AprilTag{0}".format(i), frame)

            key = cv2.waitKey(1)
            if detect_flag:
                if not locate_flag and key == ord('o') & 0xFF:
                    locate_flag = True

    except:
        traceback.print_exc()
    for cap, _ in caps:
        cap.release()


# 单相机采集
def one_camera_grub(camera_type):
    # single camera
    cap = camera_thread(camera_type)
    K_0, C_0 = read_yaml(camera_type)[1:3]
    try:
        # apriltag info
        order = "08"

        w = apriltag_size[order]["w"]
        h = apriltag_size[order]["h"]
        left = apriltag_size[order]["left"]
        right = apriltag_size[order]["right"]
        down = apriltag_size[order]["down"]
        up = apriltag_size[order]["up"]

        r, frame = cap.read()

        cv2.namedWindow("AprilTag0", cv2.WINDOW_NORMAL)
        cv2.resizeWindow("AprilTag0", 500, 300)
        cv2.moveWindow("AprilTag0", 100, 100)

        # 创建一个apriltag检测器
        at_detector = apriltag.Apriltag()
        at_detector.create_detector()

        # 预显示画面，根据apriltag显示调整顺序
        # 绘制顺序颜色，红蓝绿浅蓝
        colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (255, 255, 0)]
        while cv2.waitKey(1) != ord('q'):
            # 进行apriltag检测
            tag = at_detector.detect(frame)
            if tag.shape[0] == 4:
                for i in range(4):
                    cv2.circle(frame, tuple(tag[i]), 2, colors[i], -1)

            cv2.imshow("AprilTag0", frame)

            r, frame = cap.read()

        point_order = input(
            "Please input order:")  # should be like 1234 2341 (1 means the right_down,2 means the left_down)

        object_points_sample = np.array([
            [left + w, down, 0],  # right_down
            [left, down, 0],  # left_down
            [left, down + h, 0],  # left_top
            [left + w, down + h, 0]  # right_top
        ],

            dtype=np.float64)

        object_points_2_sample = np.array([[left + w + right, 0, 0],  # right_down
                                           [0, 0, 0],  # left_down
                                           [0, down + h + up, 0],  # left_top
                                           [left + w + right, down + h + up, 0]],  # right_top
                                          dtype=np.float64)

        object_points = object_points_sample.copy()
        object_points_2 = object_points_2_sample.copy()
        if len(point_order):
            assert len(point_order) == 4
            for i in range(4):
                object_points[point_order.find(str(i))] = object_points_sample[i]
                object_points_2[point_order.find(str(i))] = object_points_2_sample[i]

        stop_criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER,
                         30, 0.001)

        locate_flag = True

        j = 0

        if os.path.exists(CALIB_DIR):
            ts = os.listdir(CALIB_DIR)
            if len(ts):
                j = np.argmax(np.array([int(t.split('_')[1]) for t in ts])) + 1
        else:
            os.mkdir(CALIB_DIR)

        key = cv2.waitKey(1)
        while key != ord('q') & 0xFF:
            r, frame = cap.read()
            if not r:
                break
            detect_flag = True
            if locate_flag:

                tag = at_detector.detect(frame)

                if tag.shape[0] != 4:
                    # detect failure, to the next frame
                    detect_flag = False
                    break

                _, R, tvec = cv2.solvePnP(object_points, np.float64(tag).reshape(-1, 2), K_0, C_0)

                p = cv2.projectPoints(object_points_2, cv2.Rodrigues(R)[0], tvec, K_0, C_0)[0].astype(
                    np.float32).reshape(-1, 1, 2)

                gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
                p = cv2.cornerSubPix(gray, p, (5, 5), (-1, -1), stop_criteria).reshape(-1, 2)

                im = frame.copy()

                for i in range(4):
                    cv2.circle(frame, tuple(p[i].astype(int)), 4, colors[i], 2)

                cv2.imshow("AprilTag0", frame)
                key = cv2.waitKey(0)
                if key == ord('c') & 0xFF:
                    # 存储标定图片，前一张为没有标记标定点的，后一张有以便查看
                    # TODO: 文件名设定
                    cv2.imwrite(f"{CALIB_DIR}/{camera_type}_c{j}.jpg", im)

                    cv2.imwrite(f"{CALIB_DIR}/{camera_type}_c{j}_p.jpg", frame)

                    # f.write("%d\n" % j)
                    with open(f"{CALIB_DIR}/{camera_type}_{j}_corner.txt", "w") as f:
                        for ip in p:
                            f.write("%0.09f %0.09f \n" % (ip[0], ip[1]))

                    j += 1
                else:
                    print("throw")

                locate_flag = False
            else:
                cv2.imshow("AprilTag0", frame)

            if not detect_flag:
                # 如果检测失败，就直接显示图片
                cv2.putText(frame, "detect_failure", (100, 100), cv2.FONT_HERSHEY_COMPLEX, 3, (0, 0, 255), 1)
                cv2.imshow("AprilTag0", frame)

            key = cv2.waitKey(1)
            if detect_flag:
                if not locate_flag and key == ord('o') & 0xFF:
                    locate_flag = True
    except:
        traceback.print_exc()
    cap.release()
